﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.Windows.Forms;
using SharpRBM.Core;
using System.IO;
using SharpRBM.DemoTools;
using System.Diagnostics;
using System.Threading.Tasks;
using SharpRBM.Core.Enumerations;

namespace SharpRBM.SimpleLetterDemo
{
    public partial class MainForm : Form
    {
        public int reconstructedCount;

        enum StateEnum { Running, Stopping, Stopped };

        public MainForm()
        {
            InitializeComponent();
            worker.WorkerReportsProgress = true;
            worker.WorkerSupportsCancellation = true;
            worker.DoWork += new DoWorkEventHandler(worker_DoWork);
            worker.ProgressChanged += new ProgressChangedEventHandler(worker_ProgressChanged);
            worker.RunWorkerCompleted += new RunWorkerCompletedEventHandler(worker_RunWorkerCompleted);
        }

        void worker_RunWorkerCompleted(object sender, RunWorkerCompletedEventArgs e) {
            State = StateEnum.Stopped;
        }

        private StateEnum State
        {
            get
            {
                if (worker.IsBusy) {
                    if (worker.CancellationPending) {
                        return StateEnum.Stopping;
                    }
                    return StateEnum.Running;
                }
                return StateEnum.Stopped;
            }
            set
            {
                switch (value) {
                    case StateEnum.Running:
                        stopButton.Enabled = true;
                        startButton.Enabled = false;
                        break;
                    case StateEnum.Stopping:
                        stopButton.Enabled = false;
                        stopButton.Enabled = false;
                        worker.CancelAsync();
                        break;
                    case StateEnum.Stopped:
                        stopButton.Enabled = false;
                        startButton.Enabled = true;
                        break;
                }
            }
        }


        private void Form1_Shown(object sender, EventArgs e)
        {
            if (VerifyThatDataZipHasBeenUncompressed() == false)
            {
                Close();
                return;
            }

            State = StateEnum.Stopped;
            ClearBoxes();
            FindImageClassifiers();
        }

        private bool VerifyThatDataZipHasBeenUncompressed()
        {
            FileInfo fileInfo = new FileInfo(Variables.DataPath); // get a FileInfo for the file
            while( Directory.Exists( Variables.VerifyPath ) == false )
            {
                DialogResult result =
                    MessageBox.Show(
                    this,
                    string.Format(
                        "Before you can run this program, you need to uncompress the {0}\\Data.Zip to {0}\\.\nPress OK when you've uncompresed the file and Cancel if you want to quit.",
                        fileInfo.Directory),
                    "SharpRBM",
                    MessageBoxButtons.OKCancel);

                if (result == DialogResult.Cancel)
                {
                    return false;
                }
            }
            return true;
        }

        private void ClearBoxes()
        {
            inputBictureBox.Image = new Bitmap(inputBictureBox.Width, inputBictureBox.Height);
            reconstructedPictureBox.Image = new Bitmap(reconstructedPictureBox.Width, reconstructedPictureBox.Height);
        }

        private BackgroundWorker worker = new BackgroundWorker();

        private class Args {
            public int hiddenNodes;
            public List<Bitmap> bitmaps;
        }

        private void FindImageClassifiers(int hiddenNodes, int inputImages)
        {
            State = StateEnum.Running;
            ClearBoxes();

            if (Layer.UseBias)
            {
                hiddenNodes++;
            }

            // Input
            Args args = new Args();
            args.hiddenNodes = hiddenNodes;
            args.bitmaps = GetInputBitmaps(10, inputImages);

            ShowInputs(args.bitmaps);

            worker.RunWorkerAsync(args);
        }

        class Progress {
            public int generation;
            public TrainingError TrainingError;
            public double secondsForTraining;
            public List<Bitmap> FeatureDetectorBitmaps;
            public List<Bitmap> ReconstructedBitmaps;
        }

        void worker_DoWork(object sender, DoWorkEventArgs e) {
            Args args = e.Argument as Args;
            List<Layer> inputs = GetInputLayers(args.bitmaps);

            DeepBeliefNetwork deepBeliefNetwork = new DeepBeliefNetwork(inputs[0].Size, args.hiddenNodes);
            DeepBeliefNetworkTrainer trainer = new DeepBeliefNetworkTrainer(deepBeliefNetwork, deepBeliefNetwork.LayerWeights[0], inputs);

            Stopwatch stopwatch = Stopwatch.StartNew();

            // Hidden            
            for (int i = 0; i < 100000; i++) {
                Progress progress = new Progress();
                progress.generation = i;

                stopwatch.Reset();
                stopwatch.Start();

                progress.TrainingError = trainer.Train();
                progress.secondsForTraining = stopwatch.ElapsedMilliseconds / 1000.0;

                progress.FeatureDetectorBitmaps = new List<Bitmap>(GenerateFeatureDetectorBitmaps(trainer));
                progress.ReconstructedBitmaps = new List<Bitmap>(GenerateReconstructedBitmaps(trainer, inputs));

                worker.ReportProgress(0, progress);

                if (worker.CancellationPending) {
                //if (Stopping) {
                    break;
                }
            }
        }

        void worker_ProgressChanged(object sender, ProgressChangedEventArgs e) {
            //throw new NotImplementedException();
            Progress progress = e.UserState as Progress;
            //label1.Text =
            //    string.Format(
            //    "Gen {0} ({4:0.000} s): ReconstructionError={1:0.00}, DetectorError={2:0.00}, LearningRate={3:0.0000}",
            //    i,
            //    error.ReconstructionError,
            //    error.FeatureDetectorError,
            //    trainer.TrainingWeights.AdjustedLearningRate,
            //    secondsForTraining);

            label1.Text =
                string.Format(
                //"Gen {0} ({4:0.000} s): ReconstructionError={1:0.00}, DetectorError={2:0.00}, LearningRate={3:0.0000}",
                "Gen {0} ({3:0.000} s): ReconstructionError={1:0.00}, DetectorError={2:0.00}",
                progress.generation,
                progress.TrainingError.ReconstructionError,
                progress.TrainingError.FeatureDetectorError,
                //trainer.TrainingWeights.AdjustedLearningRate,
                progress.secondsForTraining);

            ShowBitmaps(featureDetectors, Enumerable.Range(0, progress.FeatureDetectorBitmaps.Count), i => progress.FeatureDetectorBitmaps[i]);
            ShowBitmaps(reconstructedPictureBox, Enumerable.Range(0, progress.ReconstructedBitmaps.Count), i => progress.ReconstructedBitmaps[i]);
        }

        private IEnumerable<Bitmap> GenerateFeatureDetectorBitmaps(DeepBeliefNetworkTrainer trainer) {
            Layer upper = new Layer(trainer.TrainingWeights.UpperLayerSize);
            Layer reconstructed = new Layer(trainer.TrainingWeights.LowerLayerSize);

            if (drawFeatureDetectorsCheckBox.Checked == false) {
                yield break;
            }

            for (int i = 0; i < upper.Size; i++) {
                upper.SetValues(0);

                if (enableBiasDuringFeatureDetectorRenderCheckBox.Checked) {
                    upper.SetBias();

                }
                else {
                    upper.Values[0] = 0;
                }

                upper.Values[i] = 1;
                trainer.LearingModule.ActivateUpperToLower(reconstructed, upper);

                if (increaseContrastCheckBox.Checked) {

                    yield return LayerBitmapHelper.CreateContrastEnhancedBitmapFromLayer(reconstructed);
                }
                else {
                    yield return LayerBitmapHelper.CreateBitmapFromLayer(reconstructed);
                }
            }
        }

        private void ShowInputs(List<Bitmap> bitmaps)
        {
            ShowBitmaps(inputBictureBox, bitmaps, bitmap => bitmap);
        }

        private static int ShowBitmaps<T>(PictureBox target, IEnumerable<T> items, Func<T, Bitmap> func)
        {            
            int count = 0;
            if (target.Height == 0 || target.Width == 0)
            {
                return 0;
            }

            Bitmap bitmap = new Bitmap(target.Width, target.Height);
            
            using (Graphics graphics = Graphics.FromImage(bitmap))
            {
                int x = 0;
                int y = 0;
                int width = LayerBitmapHelper.BitmapSize + LayerBitmapHelper.BitmapMargin;
                foreach (T item in items)
                {
                    if (x + width > target.Width)
                    {
                        x = 0;
                        y += width;

                        if (y + width > target.Height)
                        {
                            break;
                        }
                    }

                    count++;
                    graphics.DrawImageUnscaled(func(item), new Point(x, y));

                    x += width;
                }

                target.Image = bitmap;
            }

            return count;
        }

        private IEnumerable<Bitmap> GenerateReconstructedBitmaps(DeepBeliefNetworkTrainer trainer, List<Layer> layers) {

            if (drawReconstructionsCheckBox.Checked == false)
            {
                yield break;
            }

            Layer hidden = new Layer(trainer.TrainingWeights.UpperLayerSize);
            Layer reconstruction = new Layer(trainer.TrainingWeights.LowerLayerSize);
            int count = 0;
            foreach (var layer in layers) {
                trainer.LearingModule.ActivateLowerToUpper(layer, hidden, ActivationMethod.Binary);
                trainer.LearingModule.ActivateUpperToLower(reconstruction, hidden);
                yield return LayerBitmapHelper.CreateBitmapFromLayer(reconstruction);
                count++;
                if (count >= reconstructedCount) {
                    yield break;
                }
            }
        }

        public double GetErrorSum(InterLayerWeights weights, List<Layer> inputs, int upperLayerSize)
        {
            Layer testLower = new Layer(inputs[0].Size);
            Layer upper = new Layer(upperLayerSize);

            return
                inputs.Sum(layer => ComputeError(weights, layer, upper, testLower));
        }

        public double ComputeError(InterLayerWeights weights, Layer lowerLayer, Layer upperLayer, Layer testLower)
        {
            weights.ActivateLowerToUpper(lowerLayer, upperLayer);
            weights.ActivateUpperToLower(testLower, upperLayer);

            double errorSum = 0;
            Helper.Loop(lowerLayer.Size, i =>
                {
                    double error = lowerLayer.Values[i] - testLower.Values[i];
                    errorSum += error * error;
                });

            return errorSum;
        }

        public List<Bitmap> GetInputBitmaps(int number, int count)
        {
            List<Bitmap> bitmaps = new List<Bitmap>();

            for (int i = 0; i < count; i++)
            {
                for (int num = 0; num < number; num++)
                {
                    string fileName = string.Format( Variables.ImagePath, i, num );
                    Bitmap bitmap = new Bitmap(fileName);
                    bitmaps.Add(bitmap);
                }
            }
            return bitmaps;
        }

        public List<Layer> GetInputLayers(List<Bitmap> bitmaps)
        {
            return
                bitmaps
                    .Select(bitmap => LayerBitmapHelper.CreateLayerFromBitmap(bitmap))
                    .ToList();
        }

        private void MainForm_FormClosing(object sender, FormClosingEventArgs e)
        {
            worker.CancelAsync();
            
            //TODO: wait.
        }

        private void startButton_Click(object sender, EventArgs e)
        {
            FindImageClassifiers();
        }

        private void FindImageClassifiers()
        {
            Layer.UseBias = useBiasCheckBox.Checked;
            FindImageClassifiers(Convert.ToInt32(hiddenNodesTextBox.Text), Convert.ToInt32(inputImagesTextBox.Text));
        }

        private void stopButton_Click(object sender, EventArgs e)
        {
            worker.CancelAsync();            
            //TODO: wait
        }

        private void SetRecontructedCount()
        {
            reconstructedCount = (int)Math.Ceiling((double)reconstructedPictureBox.Width / (double)(LayerBitmapHelper.BitmapSize + LayerBitmapHelper.BitmapMargin));
        }

        private void MainForm_SizeChanged(object sender, EventArgs e)
        {
            SetRecontructedCount();
        }

        private void MainForm_Load(object sender, EventArgs e)
        {
            SetRecontructedCount();
        }
    }
}
